import logging
from os import listdir
from os.path import splitext
from pathlib import Path
import numpy as np,cv2
import torch
from PIL import Image
from torch.utils.data import Dataset
from dataset import data_aug
import matplotlib.pyplot as plt
from utils import config



#--------------------
def img_pad(pil_file):
    w, h = pil_file.size
    fixed_size = 1024  # 输出正方形图片的尺寸

    if h >= w:
        factor = h / float(fixed_size)
        new_w = int(w / factor)
        if new_w % 2 != 0:
            new_w -= 1
        pil_file = pil_file.resize((new_w, fixed_size))
        pad_w = int((fixed_size - new_w) / 2)
        array_file = np.array(pil_file)
        array_file = np.pad(array_file, ((0, 0), (pad_w, pad_w)), 'constant')
    else:
        factor = w / float(fixed_size)
        new_h = int(h / factor)
        if new_h % 2 != 0:
            new_h -= 1
        pil_file = pil_file.resize((fixed_size, new_h))
        pad_h = int((fixed_size - new_h) / 2)
        array_file = np.array(pil_file)
        array_file = np.pad(array_file, ((pad_h, pad_h), (0, 0)), 'constant')

    output_file = Image.fromarray(array_file)
    return output_file


class BasicDataset(Dataset):
    def __init__(self, images_dir: str, masks_dir: str,use_net=None, mask_suffix: str = ''):
        self.images_dir = Path(images_dir)
        self.masks_dir = Path(masks_dir)

        self.use_net=use_net
        self.mask_suffix = mask_suffix

        self.ids = [splitext(file)[0] for file in listdir(images_dir) if not file.startswith('.')]
        if not self.ids:
            raise RuntimeError(f'No input file found in {images_dir}, make sure you put your images there')
        logging.info(f'Creating dataset with {len(self.ids)} examples')

    def __len__(self):
        return len(self.ids)

    @classmethod
    def preprocess(cls, pil_img,config,use_net, is_mask):
        origin_w, origin_h = pil_img.size
        if use_net is "UNet":
            #----------------------------------
            #          oringin resize
            # ----------------------------------
            assert 0 < config.img_scale <= 1, 'Scale must be between 0 and 1'
            newW, newH = int(config.img_scale * origin_w), int(config.img_scale * origin_h)
            assert newW > 0 and newH > 0, 'Scale is too small, resized images would have no pixel'
            pil_img = pil_img.resize((newW, newH), resample=Image.NEAREST if is_mask else Image.BICUBIC)
            img_ndarray = np.asarray(pil_img)
            #---------------------------------------------------------
        else:
            if origin_h>origin_w:
                newH = config.max_size
                newW = int(newH * origin_w / origin_h)
                if newW % 2 != 0:
                    newW -= 1
                assert newW > 0 and newH > 0, 'Scale is too small, resized images would have no pixel'
                pil_img = pil_img.resize((newW, newH), resample=Image.NEAREST if is_mask else Image.BICUBIC)
                pad_w = int((config.max_size - newW) / 2)
                img_ndarray = np.asarray(pil_img)
                if is_mask:
                    img_ndarray = np.pad(img_ndarray, ((0, 0),(pad_w, pad_w)), 'constant')
                else:
                    #np.pad(img_ndarray, (h=(0, 0),w=(pad_h, pad_h), c=(0, 0)), 'constant')
                    img_ndarray = np.pad(img_ndarray, ((0, 0),(pad_w, pad_w), (0, 0)), 'constant')
            else:
                newW = config.max_size
                newH = int(newW * origin_h / origin_w)
                if newH%2 !=0:
                    newH -=1
                assert newW > 0 and newH > 0, 'Scale is too small, resized images would have no pixel'
                pil_img = pil_img.resize((newW, newH), resample=Image.NEAREST if is_mask else Image.BICUBIC)
                pad_h = int((config.max_size - newH) / 2)
                img_ndarray = np.asarray(pil_img)
                if is_mask:
                    img_ndarray = np.pad(img_ndarray, ((pad_h, pad_h), (0, 0)), 'constant')
                else:
                    img_ndarray = np.pad(img_ndarray, ((pad_h, pad_h), (0, 0), (0, 0)), 'constant')

        #----------------------------------
        #          visible
        # ----------------------------------
        # if not is_mask:
        #     img=cv2.cvtColor(img_ndarray,cv2.COLOR_RGB2BGR)
        #     cv2.imshow("ss",img)
        #     cv2.waitKey(0)
        #     cv2.destroyAllWindows()
        # else:
        #     # img = cv2.cvtColor(img_ndarray, cv2.COLOR_RGB2BGR)
        #     cv2.imshow("mask", img_ndarray)
        #     cv2.waitKey(0)
        #     cv2.destroyAllWindows()

        #----------------------------------------
        if img_ndarray.ndim == 2 and not is_mask:
            img_ndarray = img_ndarray[np.newaxis, ...]
        elif not is_mask:
            img_ndarray = img_ndarray.transpose((2, 0, 1))

        if not is_mask:
            img_ndarray = img_ndarray / 255

        return img_ndarray

    @classmethod
    def load(cls, filename):
        ext = splitext(filename)[1]
        if ext in ['.npz', '.npy']:
            return Image.fromarray(np.load(filename))
        elif ext in ['.pt', '.pth']:
            return Image.fromarray(torch.load(filename).numpy())
        else:
            return Image.open(filename)

    def __getitem__(self, idx):
        name = self.ids[idx]
        # mask_file = list(self.masks_dir.glob(name + self.mask_suffix + '.*'))
        mask_file = list(self.masks_dir.glob(name + '.*'))
        img_file = list(self.images_dir.glob(name + '.*'))

        assert len(mask_file) == 1, f'Either no mask or multiple masks found for the ID {name}: {mask_file}'
        assert len(img_file) == 1, f'Either no image or multiple images found for the ID {name}: {img_file}'
        mask = self.load(mask_file[0])
        img = self.load(img_file[0])  #<class 'PIL.PngImagePlugin.PngImageFile'>
        assert img.size == mask.size, \
            'Image and mask {name} should be the same size, but are {img.size} and {mask.size}'

        #----------------------------------------
        #               data bug
        #----------------------------------------
        img, mask = data_aug.Random_imgaug(np.asanyarray(img),np.asanyarray(mask),flag=np.random.randint(0,12))
        img=Image.fromarray(img.astype("uint8")).convert("RGB")
        mask=Image.fromarray(mask.get_arr())
        img = self.preprocess(img, config, self.use_net,is_mask=False)   #ndarray
        mask = self.preprocess(mask, config,self.use_net,is_mask=True)  ##ndarray

        return {
            'image': torch.as_tensor(img.copy()).float().contiguous(),
            'mask': torch.as_tensor(mask.copy()).long().contiguous()
        }


class CarvanaDataset(BasicDataset):
    def __init__(self, images_dir, masks_dir,use_net=None):
        super().__init__(images_dir, masks_dir, mask_suffix='_mask')



if __name__ == '__main__':
    pass

    # from torchvision import transforms
    # train_transform = transforms.Compose([
    #     RandomBGRRGBInverse(),
    # ])

    # da = TestPaperDataset(base_dataset, config, transforms=train_transform)
    # # print(intersect_between_two_line_segment([[0, 0], [10, 0]], [[1, -1], [1, 10]]))
    #
    # print(len(da))
    # _ = da[2]
    #
    # for i in range(len(da)):
    #     _ = da[i]
    #     print(i)
